import matplotlib.pyplot as plt
import numpy as np
from utils.data_to_color import ImColor_tp, ImColor_tp_1hour
from utils.mssql import sqlDB
# from utils.station_inter import datapro
import os
# from datetime import datetime, timedelta
from utils.interpolation import IDW, krige
from scipy.interpolate import griddata
import time

reso = 0.01
# latmin, latmax, lonmin, lonmax = 27.75, 32.38, 119.19, 123.82
# latmin, latmax, lonmin, lonmax = 29.28, 30.24, 119.99, 121.20
latmin, latmax, lonmin, lonmax = 29.25, 30.26, 119.97, 121.22

# xdim, ydim, reso = 464, 464, 0.01
xdim, ydim = int((latmax - latmin) / reso) + 1, int((lonmax - lonmin) / reso) + 1

grid_lat = np.linspace(latmax, latmin, xdim).round(2)
grid_lon = np.linspace(lonmin, lonmax, ydim).round(2)
ygrid, xgrid = np.meshgrid(grid_lon, grid_lat)

# 将插值网格数据整理

# 这里将数组转成列表
grid_lat_list = xgrid.flatten().tolist()
grid_lon_list = ygrid.flatten().tolist()

test_DB = sqlDB()
# 用一个rs变量获取数据

hourcnt = 10
radar_lhrs_data = np.zeros((hourcnt, xdim, ydim))
radar_idw_data = np.zeros((hourcnt, xdim, ydim)) * np.nan
radar_floder = "/home/gym/datasets/rain/shaoxing"

radar_files = sorted(os.listdir(radar_floder))[100:140]
for i, file in enumerate(radar_files):
    rad_file = f"{radar_floder}/{file}"
    # print(i, rad_file)
    rdata = np.load(rad_file).round(2)[212: 314, 78:204]
    # plt.imshow(ImColor_tp_1hour(rdata))  #
    # plt.title(f'radar_data')
    # plt.show()
    if i < hourcnt:
        radar_lhrs_data[i] = rdata
    else:
        rzero = radar_lhrs_data.copy()[1:]
        radar_lhrs_data[:9] = rzero
        radar_lhrs_data[9] = rdata

        rdar_mean = np.mean(radar_lhrs_data, axis=0).round(2)
        # radar_utc_time = datetime.strptime(os.path.basename(rad_file)[:12], '%Y%m%d%H%M')
        # station_bj_time = (radar_utc_time + timedelta(hours=8)).strftime('%Y%m%d%H%M')

        station_bj_time = os.path.basename(rad_file)[:12]
        # ftime = str(int(ftime0 - ftime0 % 5 + int((ftime0 % 5 + 2.5) / 5) * 5))
        # print(station_bj_time)
        results = test_DB.select(station_bj_time, timerange=(-60, 0))
        station_data = np.zeros_like(rdar_mean)
        know_lon, know_lat, know_z = [], [], []
        for stdata in results:
            StationNum, EEEee, NNnn, Precipitation = stdata
            if StationNum is None or EEEee is None or NNnn is None or Precipitation is None:
                continue

            Precipitation /= 10.

            lon = float(EEEee) / 100
            lat = float(NNnn) / 100
            x = int((latmax - lat) / reso + 0.5)
            y = int((lon - lonmin) / reso + 0.5)

            # station_data[x, y] = Precipitation
            r = 1
            station_data[x - r:x + r, y - r:y + r] = Precipitation

            rdar_d = rdar_mean[x, y]
            dz = round(rdar_d - Precipitation, 2)

            # latl.append(int((float(NNnn) / 100 - latmin) / reso + 0.5))
            know_lon.append(lon)
            know_lat.append(lat)
            know_z.append(dz)

        points = tuple((np.array(know_lon), np.array(know_lat)))
        values = np.array(know_z)

        # station_1min_reults = test_DB.select(station_bj_time, timerange=(-3, 3))
        # station_data_1min = np.zeros_like(rdar_mean)
        # for stdata in station_1min_reults:
        #     StationNum, EEEee, NNnn, Precipitation = stdata
        #     if StationNum is None or EEEee is None or NNnn is None or Precipitation is None:
        #         continue
        #
        #     Precipitation /= 10.
        #
        #     lon = float(EEEee) / 100
        #     lat = float(NNnn) / 100
        #     x = int((latmax - lat) / reso + 0.5)
        #     y = int((lon - lonmin) / reso + 0.5)
        #     xarr.append(x)
        #     yarr.append(y)
        #     # station_data[x, y] = Precipitation
        #     r = 1
        #     station_data_1min[x - r:x + r, y - r:y + r] = Precipitation


        st = time.time()
        _idw = IDW(know_lon, know_lat, know_z, grid_lon_list, grid_lat_list)
        # print("IDW TIME COST", time.time() - st)
        pm_idw = np.array(_idw).reshape((xdim, ydim, 3))[..., 2]
        pm_idw = np.minimum(np.nan_to_num(pm_idw), rdar_mean)

        st1 = time.time()
        pm_krige = krige(know_lon, know_lat, know_z, grid_lon, grid_lat)
        # print("Krige TIME COST", time.time() - st1)
        pm_krige = np.minimum(np.nan_to_num(pm_krige), rdar_mean)

        xidxs = ((latmax - np.array(know_lat)) / reso + 0.5).astype(np.int32)
        yidxs = ((np.array(know_lon) - lonmin) / reso + 0.5).astype(np.int32)
        # real_rain = griddata(points, station_data[(xidxs, yidxs)], (ygrid, xgrid), method='cubic', fill_value=0,
        #                          rescale=True)
        # real_rain = krige(know_lon, know_lat, station_data[(xidxs, yidxs)], grid_lon, grid_lat)
        # _rain = IDW(know_lon, know_lat, station_data[(xidxs, yidxs)], grid_lon_list, grid_lat_list)
        # real_rain = np.array(_rain).reshape((xdim, ydim, 3))[..., 2]
        # real_rain = np.maximum(np.nan_to_num(real_rain), 0)
        real_rain = np.maximum(0, np.nan_to_num(rdar_mean - pm_idw).round(1))

        idw_rain = np.maximum(0, np.nan_to_num(rdata - pm_idw).round(1))
        krige_rain = np.maximum(0, np.nan_to_num(rdata - pm_krige).round(1))




        krige_rain10 = radar_idw_data.copy()
        idw_idx = i - 9
        if idw_idx < 9:
            radar_idw_data[idw_idx] = idw_rain
            krige_rain10[idw_idx] = krige_rain
        else:
            _rzero = np.zeros_like(radar_lhrs_data)
            _rzero[:9] = radar_idw_data.copy()[1:]
            _rzero[9] = idw_rain
            radar_idw_data = _rzero

            rzero = np.zeros_like(krige_rain10)
            rzero[:9] = krige_rain10.copy()[1:]
            rzero[9] = krige_rain
            krige_rain10 = rzero

        idw_rain_mean = np.nanmean(radar_idw_data, axis=0).round(1)
        krige_rain_mean = np.nanmean(krige_rain10, axis=0).round(1)
        # mse = np.square(station_data[(xidxs, yidxs)] - idw_rain_mean[(xidxs, yidxs)]).mean()
        rdar_mse = np.square(station_data - rdar_mean)[(xidxs, yidxs)].mean()
        idw_mse = np.square(station_data - idw_rain_mean)[(xidxs, yidxs)].mean()
        krige_mse = np.square(station_data - krige_rain_mean)[(xidxs, yidxs)].mean()

        print(f"Rdar_mse:{rdar_mse:.4f}, IDW_mse:{idw_mse:.4f}, Krige_mse:{krige_mse:.4f}", )

        plt.suptitle(f"{station_bj_time}")

        plt.subplot(231)
        plt.imshow(ImColor_tp_1hour(rdata))  #
        plt.title(f'radar_data')

        plt.subplot(232)
        plt.imshow(ImColor_tp_1hour(rdar_mean))  #
        plt.title(f'rdar_mean')
        #
        plt.subplot(233)
        plt.imshow(ImColor_tp_1hour(station_data))  # [214: 310, 80:202]
        plt.title(f'station_data')

        plt.subplot(234)
        plt.imshow(ImColor_tp_1hour(idw_rain_mean))  # [214: 310, 80:202]
        plt.title(f'idw_rain')

        plt.subplot(235)
        # rain_cubic[radar_data[i] < 0.1] = 0
        plt.imshow(ImColor_tp_1hour(real_rain))  # [214: 310, 80:202]
        plt.title(f'real_rain')

        # plt.subplot(234)
        # rain_cubic = radar_data[i] - cubic
        # # rain_cubic[radar_data[i] < 0.1] = 0
        # plt.imshow(ImColor_tp_1hour(rain_cubic))  # [214: 310, 80:202]
        # plt.title(f'inter_cubic')

        plt.subplot(236)
        plt.imshow(ImColor_tp_1hour(krige_rain_mean))  # [214: 310, 80:202]
        plt.title(f'krige_rain')

        # plt.subplot(236)
        # plt.imshow(ImColor_tp_1hour(idw_rain_mean))  # [214: 310, 80:202]
        # plt.title(f'idw_rain_mean')
        #
        # plt.subplot(133)
        # plt.imshow(ImColor_tp_12mins(rs_data.copy()[214: 310, 80:202] * 2))  #
        # plt.title(f'rs_data {station_bj_time}')

        plt.gcf().set_size_inches(16, 16)
        plt.tight_layout()
        plt.show()
        result_file = f"results/{station_bj_time}.jpg"
        # plt.savefig(result_file)
        print(i, result_file)
# exit()

test_DB.close()
